// StepTrack.C
#include <iostream>
#include <unistd.h>

#include "LOCA.H"
#include "LOCA_Epetra.H"

#include "Epetra_SerialComm.h"
#include "Epetra_Map.h"
#include "Epetra_Vector.h"
#include "Epetra_RowMatrix.h"
#include "Epetra_CrsMatrix.h"
#include "Epetra_Map.h"
#include "Epetra_LinearProblem.h"
#include "AztecOO.h"

#include "Global.H"
#include "StepWaveProb.H"              
#include "StepProbInterface.H"

#define ARGS1    "m n p0 p1 p2 a1 a2 a3 contParam stepSize numSteps"
#define NUMARGS1 11 
#define ARGS2    "solfile contParam stepSize numSteps"
#define NUMARGS2 4


using namespace std;


int main(int argc, char *argv[])
{
  try
  {
    StepWaveProb *prob;
    char contParam[1];
    double stepSize;
    int numSteps;

    if (argc == NUMARGS1+1)
    {
      int i = 1;
      int m = atoi(argv[i++]);
      int n = atoi(argv[i++]);
      double p0 = atof(argv[i++]); if (p0 > 0) { cerr << "Expected negative flux, exiting" << endl; exit(1); }
      double p1 = atof(argv[i++]);
      double p2 = atof(argv[i++]);
      double a1 = atof(argv[i++]);
      double a2 = atof(argv[i++]);
      double a3 = atof(argv[i++]);
      sprintf(contParam, "%c", toupper(argv[i++][0]));
      stepSize = atof(argv[i++]);
      numSteps = atoi(argv[i++]);

      prob = new StepWaveProb(m, n, p0, p1, p2, a1, a2, a3);
    }
    else if (argc == NUMARGS2+1)
    {
      int i=1;
      char *solfile = argv[i++];
      sprintf(contParam, "%c", toupper(argv[i++][0]));
      stepSize = atof(argv[i++]);
      numSteps = atoi(argv[i++]);

      prob = new StepWaveProb(solfile);
    }
    else { fprintf(stderr, "usage:\n  %s %s\n  %1$s %3$s\n", argv[0], ARGS1, ARGS2); exit(1); }


    double contParamVal;
    if (strcmp(contParam,"Q")==0)
      contParamVal = prob->getQ();
    else if (strcmp(contParam,"P1")==0)
      contParamVal = prob->getP1();
    else if (strcmp(contParam,"P2")==0)
      contParamVal = prob->getP2();
    else if (strcmp(contParam,"A1")==0)
      contParamVal = prob->getA1();
    else if (strcmp(contParam,"A2")==0)
      contParamVal = prob->getA2();
    else if (strcmp(contParam,"A3")==0)
      contParamVal = prob->getA3();
    

    // Get the solution vector from the problem
    Epetra_Vector& soln = prob->getInitial();


    // Begin LOCA Solver ************************************

    // Create parameter list
    NOX::Parameter::List paramList;

    // Create LOCA sublist
    NOX::Parameter::List& locaParamsList = paramList.sublist("LOCA");

    // Create the stepper sublist and set the stepper parameters
    NOX::Parameter::List& locaStepperList = locaParamsList.sublist("Stepper");
    locaStepperList.setParameter("Continuation Method", "Householder Arc Length");
    locaStepperList.setParameter("Continuation Parameter", contParam);
    locaStepperList.setParameter("Initial Value", contParamVal);
    locaStepperList.setParameter("Max Value", 5000.0);
    locaStepperList.setParameter("Min Value", 0.0);
    locaStepperList.setParameter("Max Steps", numSteps);
    locaStepperList.setParameter("Max Nonlinear Iterations", 10);
    locaStepperList.setParameter("Enable Arc Length Scaling", true);
    locaStepperList.setParameter("Goal Arc Length Parameter Contribution", 0.5);
    locaStepperList.setParameter("Max Arc Length Parameter Contribution", 0.7);
    locaStepperList.setParameter("Initial Scale Factor", 1.0);
    locaStepperList.setParameter("Min Scale Factor", 1.0e-8);
    locaStepperList.setParameter("Enable Tangent Factor Step Size Scaling", true);
    locaStepperList.setParameter("Min Tangent Factor", -1.8);
    locaStepperList.setParameter("Tangent Factor Exponent", 1.0);

    // Create bifurcation sublist
    NOX::Parameter::List& bifurcationList = locaParamsList.sublist("Bifurcation");
    bifurcationList.setParameter("Method", "None");

    // Create Anasazi Eigensolver sublist (needs --with-loca-anasazi)
    locaStepperList.setParameter("Compute Eigenvalues", true);
    NOX::Parameter::List& aList = locaStepperList.sublist("Anasazi");
    aList.setParameter("Block Size", 1);
    aList.setParameter("Arnoldi Size", 30);
    aList.setParameter("NEV", 4);
    aList.setParameter("Save Eigenvectors", 3);
    aList.setParameter("Tol", 2.0e-5);
    aList.setParameter("Convergence Check", 1);
    aList.setParameter("Restarts", 2);
    aList.setParameter("Frequency", 1);
    aList.setParameter("Debug Level", 0);

    // Create predictor sublist
    NOX::Parameter::List& predictorList = locaParamsList.sublist("Predictor");
    predictorList.setParameter("Method", "Tangent");

    // Create step size sublist
    NOX::Parameter::List& stepSizeList = locaParamsList.sublist("Step Size");
    stepSizeList.setParameter("Initial Step Size", stepSize);
    stepSizeList.setParameter("Method", "Constant");
    stepSizeList.setParameter("Min Step Size", 1.0e-5);
    stepSizeList.setParameter("Max Step Size", 10.0);
    stepSizeList.setParameter("Aggressiveness", 0.5);
    stepSizeList.setParameter("Failed Step Reduction Factor", 0.5);

    // Set the LOCA Utilities
    NOX::Parameter::List& locaUtilsList = locaParamsList.sublist("Utilities");
    //locaUtilsList.setParameter("MyPID", MyPID);
    locaUtilsList.setParameter("Output Information", 
        LOCA::Utils::Warning +
        LOCA::Utils::StepperIteration +
        LOCA::Utils::StepperDetails +
        LOCA::Utils::Solver +
        LOCA::Utils::SolverDetails +
        LOCA::Utils::Parameters);

    // Create the "Solver" parameters sublist to be used with NOX Solvers
    NOX::Parameter::List& nlParams = paramList.sublist("NOX");
    nlParams.setParameter("Nonlinear Solver", "Line Search Based");

    // Create the NOX printing parameter list
    NOX::Parameter::List& nlPrintParams = nlParams.sublist("Printing");
    //nlPrintParams.setParameter("MyPID", MyPID); 
    nlPrintParams.setParameter("Output Information", NOX::Utils::Warning);
    nlPrintParams.setParameter("Output Information",
        NOX::Utils::OuterIteration + 
        NOX::Utils::OuterIterationStatusTest + 
        NOX::Utils::InnerIteration +
        NOX::Utils::Parameters + 
        NOX::Utils::Details + 
        NOX::Utils::Warning);

    // Create the "Line Search" sublist for the "Line Search Based" solver
    NOX::Parameter::List& searchParams = nlParams.sublist("Line Search");
    searchParams.setParameter("Method", "Full Step");
    searchParams.setParameter("Max Iters", 20);
    searchParams.setParameter("Default Step", 1.0000);
    searchParams.setParameter("Recovery Step", 0.0001);
    searchParams.setParameter("Minimum Step", 0.0001);

    // Create the "Direction" sublist for the "Line Search Based" solver
    NOX::Parameter::List& dirParams = nlParams.sublist("Direction");
    NOX::Parameter::List& newParams = dirParams.sublist("Newton");
    dirParams.setParameter("Method", "Newton");
    newParams.setParameter("Forcing Term Method", "Constant");

    // Create the "Linear Solver" sublist for the "Direction" sublist
    NOX::Parameter::List& lsParams = newParams.sublist("Linear Solver");
    lsParams.setParameter("Aztec Solver", "GMRES");  
    lsParams.setParameter("Max Iterations", 100);  
    lsParams.setParameter("Tolerance", 1e-5);
    lsParams.setParameter("Output Frequency", 50);
    lsParams.setParameter("Scaling", "None");
    lsParams.setParameter("Preconditioner", "AztecOO");
    lsParams.setParameter("Aztec Preconditioner", "ilut"); 

    // Create and initialize the parameter vector
    LOCA::ParameterVector pVector;
    pVector.addParameter(contParam, contParamVal);
    
    // Create the interface between the test problem and the nonlinear solver
    // This is created by the user using inheritance of the abstract base class:
    // NLS_PetraGroupInterface
    StepProbInterface interface(*prob, contParam);

    // Create the Epetra_RowMatrix for the Jacobian/Preconditioner
    Epetra_CrsMatrix & J = prob->getJacobian();

    // Create the linear systems
    NOX::EpetraNew::LinearSystemAztecOO
      linsys(nlPrintParams, lsParams, interface, interface, J, soln);

    // Create the loca vector
    NOX::Epetra::Vector locaSoln(soln);

    // Create the Group
    LOCA::EpetraNew::Group grp(nlPrintParams, interface, locaSoln, linsys, pVector);
    grp.computeF();

    // Create the Solver convergence test
    NOX::StatusTest::NormF wrms(1e-5);
    NOX::StatusTest::MaxIters maxiters(searchParams.getParameter("Max Iters", 20));
    NOX::StatusTest::Combo combo(NOX::StatusTest::Combo::OR);
    combo.addStatusTest(wrms);
    combo.addStatusTest(maxiters);

    // Create the stepper  
    LOCA::Stepper stepper(grp, combo, paramList);

    // Perform continuation run
    stepper.run();

    // Output the parameter list
    if (LOCA::Utils::doPrint(LOCA::Utils::Parameters))
    {
      cout << "*********************************************************" << endl
           << "Continuation Finished. Final Parameters:" << endl;
      stepper.getParameterList().print(cout);
      cout << endl;
    }


    delete prob;
  }
  catch (char *s) { cout << s << endl; }
  catch (...) { cout << "Caught unknown exception" << endl; }
}
